import base64
import logging
import shutil
import threading
from functools import wraps
from io import BytesIO
from server import PromptServer
from aiohttp import web
import folder_paths
from .src.config import *
from .src.database import get_image_path_by_id, is_video_exist, get_pexels_video_count
from .src.models import DatabaseSession, DatabaseSessionPexelsVideo
from .src.process_assets import match_text_and_image, process_image, process_text
from .src.scan import Scanner
from .src.search import (
    clean_cache,
    search_image_by_image,
    search_image_by_text_path_time,
    search_video_by_image,
    search_video_by_text_path_time,
    search_pexels_video_by_text,
)
from .src.utils import crop_video, get_hash, resize_image_with_aspect_ratio
import json
logging.basicConfig(
    level=LOG_LEVEL, format="%(asctime)s %(name)s %(levelname)s %(message)s"
)
logger = logging.getLogger(__name__)

scanner = Scanner()

@PromptServer.instance.routes.get("/search/app")
async def app(request):
    basePath = folder_paths.folder_names_and_paths['custom_nodes'][0][0]
    htmlPtah = os.path.join(basePath, 'ComfyUI_MaterialSearch', 'static','index.html')
    return web.FileResponse(htmlPtah)

@PromptServer.instance.routes.get("/search/api/scan")
async def scan(request):
    """开始扫描"""
    global scanner
    if not scanner.is_scanning:
        scan_thread = threading.Thread(target=scanner.scan, args=(False,))
        scan_thread.start()
        return web.Response(text=json.dumps({"status": "start scanning"}), content_type='application/json')
    return web.Response(text=json.dumps({"status": "already scanning"}), content_type='application/json')


@PromptServer.instance.routes.get("/search/api/status")
async def status(request):
    """状态"""
    global scanner
    result = scanner.get_status()
    with DatabaseSessionPexelsVideo() as session:
        result["total_pexels_videos"] = get_pexels_video_count(session)
    return web.Response(text=json.dumps(result), content_type='application/json')

@PromptServer.instance.routes.get("/search/api/clean_cache")
async def clean_cache_(request):
    """清理缓存"""
    clean_cache()
    return web.Response(text=json.dumps({"status": "clean cache"}), content_type='application/json')

@PromptServer.instance.routes.post("/search/api/match")
async def match(request):
    """
    匹配文字对应的素材
    :return: json格式的素材信息列表
    """
    #post = await request.post()
    data =  await request.json()
    top_n = int(data["top_n"])
    search_type = data["search_type"]
    positive_threshold = data["positive_threshold"]
    negative_threshold = data["negative_threshold"]
    image_threshold = data["image_threshold"]
    img_id = data["img_id"]
    path = data["path"]
    start_time = data["start_time"]
    end_time = data["end_time"]
    upload_file_path = data["upload_file_path"] if 'upload_file_path' in data else None #image = post.get("upload_file_path") #session.get('upload_file_path', '')
    if search_type in (1, 3, 4):
        if not upload_file_path or not os.path.exists(upload_file_path):
            return "你没有上传文件！", 400
    logger.debug(data)
    # 进行匹配
    if search_type == 0:  # 文字搜图
        results = search_image_by_text_path_time(data["positive"], data["negative"], positive_threshold, negative_threshold,
                                                 path, start_time, end_time)
    elif search_type == 1:  # 以图搜图
        results = search_image_by_image(upload_file_path, image_threshold)
    elif search_type == 2:  # 文字搜视频
        results = search_video_by_text_path_time(data["positive"], data["negative"], positive_threshold, negative_threshold,
                                                 path, start_time, end_time)
    elif search_type == 3:  # 以图搜视频
        results = search_video_by_image(upload_file_path, image_threshold)
    elif search_type == 4:  # 图文相似度匹配
        score = match_text_and_image(process_text(data["text"]), process_image(upload_file_path)) * 100
        return web.Response(text=json.dumps({"score": "%.2f" % score}), content_type='application/json')
    elif search_type == 5:  # 以图搜图(图片是数据库中的)
        results = search_image_by_image(img_id, image_threshold)
    elif search_type == 6:  # 以图搜视频(图片是数据库中的)
        results = search_video_by_image(img_id, image_threshold)
    elif search_type == 9:  # 文字搜pexels视频
        results = search_pexels_video_by_text(data["positive"], positive_threshold)
    else:  # 空
        logger.warning(f"search_type不正确：{search_type}")
        return web.Response(status=400)
    return web.Response(text=json.dumps(results[:top_n]), content_type='application/json')

@PromptServer.instance.routes.get("/search/api/get_image/{image_id}")
async def api_get_image(request):
    """
    读取图片
    :param image_id: int, 图片在数据库中的id
    :return: 图片文件
    """
    image_id = request.match_info.get("image_id", None)
    if image_id is None:
        return web.Response(status=404)
    image_id=int(image_id)
    with DatabaseSession() as session:
        path = get_image_path_by_id(session, image_id)
        logger.debug(path)
    thumbnail=''
    if "thumbnail" in request.rel_url.query:
        thumbnail=request.rel_url.query['thumbnail']
    # 静态图片压缩返回
    if thumbnail and os.path.splitext(path)[-1] != "gif":
        image = resize_image_with_aspect_ratio(path, (640, 480), convert_rgb=True)
        image_io = BytesIO()
        image.save(image_io, 'JPEG', quality=60)
        image_io.seek(0)
        # 返回文件
        return web.Response(body=image_io)
    return web.FileResponse(path)


@PromptServer.instance.routes.get("/search/api/get_video/{video_path}")
async def api_get_video(request):
    """
    读取视频
    :param video_path: string, 经过base64.urlsafe_b64encode的字符串，解码后可以得到视频在服务器上的绝对路径
    :return: 视频文件
    """
    video_path = request.match_info.get("video_path", None)
    if video_path is None:
        return web.Response(status=404)
    path = base64.urlsafe_b64decode(video_path).decode()
    logger.debug(path)
    with DatabaseSession() as session:
        if not is_video_exist(session, path):  # 如果路径不在数据库中，则返回404，防止任意文件读取攻击
            return web.Response(status=404)
    return web.FileResponse(path)


@PromptServer.instance.routes.get("/search/api/download_video_clip/{video_path}/{start_time}/{end_time}")
async def api_download_video_clip(request):
    """
    下载视频片段
    :param video_path: string, 经过base64.urlsafe_b64encode的字符串，解码后可以得到视频在服务器上的绝对路径
    :param start_time: int, 视频开始秒数
    :param end_time: int, 视频结束秒数
    :return: 视频文件
    """
    video_path=request.match_info.get("video_path", None)
    start_time=request.match_info.get("start_time", None)
    end_time=request.match_info.get("end_time", None)
    start_time, end_time = int(start_time), int(end_time)
    path = base64.urlsafe_b64decode(video_path).decode()
    logger.debug(path)
    with DatabaseSession() as session:
        if not is_video_exist(session, path):  # 如果路径不在数据库中，则返回404，防止任意文件读取攻击
            return web.Response(status=404)
    # 根据VIDEO_EXTENSION_LENGTH调整时长
    start_time -= VIDEO_EXTENSION_LENGTH
    end_time += VIDEO_EXTENSION_LENGTH
    if start_time < 0:
        start_time = 0
    # 调用ffmpeg截取视频片段
    output_path = f"{TEMP_PATH}/video_clips/{start_time}_{end_time}_" + os.path.basename(path)
    if not os.path.exists(output_path):  # 如果存在说明已经剪过，直接返回，如果不存在则剪
        crop_video(path, output_path, start_time, end_time)
    return web.FileResponse(output_path)

@PromptServer.instance.routes.post("/search/api/upload")
async def upload(request):
    """
    上传文件。首先删除旧的文件，保存新文件，计算hash，重命名文件。
    :return: 200
    """
    # 保存文件
    post = await request.post()
    image = post.get("file")
    filename = image.filename
    if not filename:
        return web.Response(status=400)
    upload_file_path = os.path.join(f"{TEMP_PATH}/upload",filename)
    with open(upload_file_path, "wb") as f:
        f.write(image.file.read())
    logger.debug(f"upload_file_path: {upload_file_path}")
    return web.Response(text=json.dumps({"upload_file_path":upload_file_path}), content_type='application/json')


def init():
    """
    清理和创建临时文件夹，初始化扫描线程（包括数据库初始化），根据AUTO_SCAN决定是否开启自动扫描线程
    """
    global scanner
    # 删除上传目录中所有文件
    shutil.rmtree(f'{TEMP_PATH}/upload', ignore_errors=True)
    os.makedirs(f'{TEMP_PATH}/upload')
    shutil.rmtree(f'{TEMP_PATH}/video_clips', ignore_errors=True)
    os.makedirs(f'{TEMP_PATH}/video_clips')
    # 初始化扫描线程
    scanner.init()
    if AUTO_SCAN:
        auto_scan_thread = threading.Thread(target=scanner.auto_scan, args=())
        auto_scan_thread.start()

init()

NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
file_directory = os.path.dirname(os.path.abspath(__file__))
WEB_DIRECTORY=file_directory+"/static/init"

PromptServer.instance.app.router.add_static("/search/static", file_directory+"/static")
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS","WEB_DIRECTORY"]